#include "mm.h"
#include "kernel.h"
#include "assert.h"

static void get_indexes(uint64_t addr, uint64_t *pml4_index, uint64_t *pdpt_index, uint64_t *pdt_index)
{
    *pml4_index = (addr >> PML4_SHIFT) & INDEX_MASK;
    *pdpt_index = (addr >> PDPT_SHIFT) & INDEX_MASK;
    *pdt_index = (addr >> PDT_SHIFT) & INDEX_MASK;
}

static uint64_t get_2mb_page(uint64_t addr)
{
    return addr & 0xffffffffffe00000;
}

void physics_map_virtual_addr_2m(uint32_t addr, uint32_t phy_addr_base)
{
    assert(NULL != addr);
    assert(addr < 0xffffffff);

    uint64_t pml4_index, pdpt_index, pdt_index;
    get_indexes(addr, &pml4_index, &pdpt_index, &pdt_index);

    printk("PML4 index:%d\n", pml4_index);
    printk("PDPT index:%d\n", pdpt_index);
    printk("PDT index: %d\n", pdt_index);

    uint64_t* PML4_addr = (uint64_t*) FOUR_LEVEL_HEAD_TABLE_ADDR;
    uint64_t pdpt_descriptor = *(PML4_addr + pml4_index);
    if (0 == pdpt_descriptor)
    {
        int64 p = (int64) get_free_page();
        *(PML4_addr + pml4_index) = p | 3;
        pdpt_descriptor = p | 3;
    }
    
    printk("PML4:%d->0x%x\n", pml4_index, pdpt_descriptor);

    uint64_t* pdpt_addr = (uint64_t*) (pdpt_descriptor >> 12 << 12);
    uint64_t pdt_descriptor = (uint64_t) *(pdpt_addr + pdpt_index);
    if (0 == pdt_descriptor)
    {
        int64 p = (int64) get_free_page();
        *(pdpt_addr + pdpt_index) = p | 3;
        pdt_descriptor = p | 3;
    }
    printk("pdpt: %d->0x%x\n", pdpt_index, pdt_descriptor);

    uint64_t* pdt_addr = (uint64_t*) (pdt_descriptor >> 12 << 12);
    *(pdt_addr + pdt_index) = phy_addr_base >> 21 << 21 | 0x83;
}

void physics_addr_map_virtual_addr_2m_2(uint64_t addr, uint32_t phy_addr_base, uint16_t pdt_flag, uint16_t pdpt_flag)
{
    assert(NULL != addr);
    assert(addr < 0xffffffff);

    uint64_t pml4_index, pdpt_index, pdt_index;
    get_indexes(addr, &pml4_index, &pdpt_index, &pdt_index);

    printk("PML4 index: %d\n", pml4_index);
    printk("PDPT index: %d\n", pdpt_index);
    printk("PDT index: %d\n", pdt_index);

    uint64_t* PML4_addr = (uint64_t) FOUR_LEVEL_HEAD_TABLE_ADDR;
    uint64_t pdpt_descriptor = *(PML4_addr + pml4_index);
    if (0 == pdpt_descriptor)
    {
        int64 p = (int64) get_free_page();
        *(PML4_addr + pml4_index) = p | 3;
        pdpt_descriptor = p | 3;
    }

    printk("PML4: %d->0x%x\n", pml4_index, pdpt_descriptor);

    uint64_t* pdpt_addr = (uint64_t*) (pdpt_descriptor >> 12 << 12);
    uint64_t pdt_descriptor = (uint64_t) *(pdpt_addr + pdpt_index);

    if (0 == pdt_descriptor)
    {
        // printk("pdt_descriptor is 0, need alloc new page\n");
        int64 p = (int64) get_free_page();
        *(pdpt_addr + pdpt_index) = p | pdpt_flag;
        pdt_descriptor = p | pdpt_flag;
    }

    printk("pdpt: %d->0x%x\n", pdpt_index, pdt_descriptor);
    
    uint64_t* pdt_addr = (uint64_t*) (pdt_descriptor >> 12 << 12);
    *(pdt_addr + pdt_index) = phy_addr_base >> 21 << 21 | pdt_flag;
}